package tcpserver

import (
	"context"
	"fmt"
	"net"
	"nprobe/core"
	"nprobe/core/agent"
	"nprobe/core/aggregate"
	"syscall"
)

// TCPServerAgent icmp echo
type TCPServerAgent struct {
	agent.BaseAgent
	TCPListeners []*net.TCPListener
}

func NewTCPServerAgent(name string, task *core.ProbeTask) (*TCPServerAgent, error) {
	tcpConns, err := listenTCPs(task)
	if err != nil {
		fmt.Println("nprobe listen tcp error: ", err)
		return nil, err
	}
	salts := core.NewSalts(task.MsgLen)

	ctx, cancelFunc := context.WithCancel(context.Background())
	receiverGroup := NewTCPServerReceiverGroup(ctx, name, task, salts)

	aggregator := aggregate.NewRingBufAggregator(ctx, name, task, nil, receiverGroup)

	return &TCPServerAgent{
		BaseAgent: agent.BaseAgent{
			Name:          name,
			Task:          task,
			SenderGroup:   nil,
			ReceiverGroup: receiverGroup,
			Aggregator:    aggregator,

			Ctx:        ctx,
			CancelFunc: cancelFunc,
		},
		TCPListeners: tcpConns,
	}, nil
}

func listenTCPs(task *core.ProbeTask) ([]*net.TCPListener, error) {
	var tcpConns []*net.TCPListener
	for port := task.DstPortMin; port <= task.DstPortMax; port++ {
		listener, err := net.ListenTCP("tcp4", &net.TCPAddr{Port: int(port)})
		if err != nil {
			for _, tcpConn := range tcpConns {
				tcpConn.Close()
			}
			return nil, err
		}

		rawConn, err := listener.SyscallConn()
		if err != nil {
			for _, tcpConn := range tcpConns {
				tcpConn.Close()
			}
			return nil, err
		}

		err = rawConn.Control(func(fd uintptr) {
			// 设置SO_REUSEADDR选项
			syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1)
		})

		tcpConns = append(tcpConns, listener)
	}
	return tcpConns, nil
}

func (t *TCPServerAgent) StopAgent() error {
	if t != nil {
		t.BaseAgent.StopAgent()

		for _, conn := range t.TCPListeners {
			conn.Close()
		}
	}
	return nil
}
